# 导入库
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
from matplotlib import pyplot as plt
import seaborn as sns
import statsmodels.api as sm
import numpy as np

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8mb4'
}

# 创建数据库引擎
engine = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}")
conn = engine.connect()
chunk_size = 10000
# 从数据库查询数据
df = pd.read_sql_query("""
                        SELECT d.*, m.buy_lg_vol, m.sell_lg_vol, m.buy_elg_vol, m.sell_elg_vol, m.net_mf_vol, i.vol as i_vol, i.closes as i_closes 
                        FROM date_1 d 
                        join moneyflows m on d.ts_code = m.ts_code and d.trade_date = m.trade_date
                        LEFT JOIN index_daily i on d.trade_date = i.trade_date and i.ts_code = '399001.SZ'
                        WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' and d.ts_code = '002229.SZ'
                        """,
                        conn,
                        chunksize=chunk_size
                        )
df1 = pd.concat(df, ignore_index=True)
# 增加一列，股票日涨跌幅
df1['zd_pct_chg'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
# 处理缺失值
df1 = df1.dropna(subset=['zd_pct_chg']).reset_index(drop=True)

numeric_cols = df1.select_dtypes(include=['number']).columns.to_list()
# 选择需要进行主成分分析的自变量
X = df1[['vol', 'amount', 'buy_lg_vol', 'sell_lg_vol', 'buy_elg_vol', 'sell_elg_vol']]
# 计算特征值和特征向量
eigenvalues, eigenvectors = np.linalg.eig(np.cov(X, rowvar=False))
print('累计共享率为：', round(eigenvalues[:5].sum() / eigenvalues.sum(), 2) * 100, '%')

# 选择要保留的主成分个数
n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]

# 计算主成分
principal_components = np.dot(X, top_eigenvectors)
# 将主成分添加到原始数据中
data_pca = pd.concat([df1, pd.DataFrame(principal_components, columns=[f'PC{i+1}' for i in range(n_components)])], axis=1)
# 添加常数项以确保数据类型正确
X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pca = sm.add_constant(X_pca)
# 确保y索引与X_pca一致
y = data_pca['zd_pct_chg'].copy()
# 构建回归模型
print(X_pca)
print(y)
model = sm.OLS(y, X_pca)
# 拟合模型
results = model.fit()
# 输出结果
print('回归模型结果：')
print(results.summary())
# 选取PC3、PC4、PC5作为新的自变量
X_pca_selected = data_pca[['PC3', 'PC4', 'PC5']]
X_pca_selected.columns = ['PC3', 'PC4', 'PC5']
# 添加常数项
X_pca_selected = sm.add_constant(X_pca_selected)
# 因变量
y = data_pca['zd_pct_chg'].copy()
# 构建回归模型
model = sm.OLS(y, X_pca_selected)
# 拟合模型
results = model.fit()
# 输出结果
print('回归模型结果：')
print(results.summary())
X_pca_selected = data_pca[['PC3', 'PC4', 'PC5']]
X_pca_selected.columns = ['PC3', 'PC4', 'PC5']
# 绘制散点图
fig, axes = plt.subplots(1, 4, figsize=(15, 5))
for i, col in enumerate(X_pca_selected.columns):
    axes[i].scatter(X_pca_selected[col], y, s=50, alpha=0.7)
    axes[i].set_xlabel(col)
    axes[i].set_ylabel('zd_pct_chg')
    axes[i].set_title(f'{col}')
plt.tight_layout()
plt.show()

# 计算相关系数
for k in range(5):
    string_y = f'CP{k+1} = '
    i = eigenvectors[k]
    for j in range(len(i)):
        if i[j] > 0:
            string_y = string_y + f' {round(i[j], 2)} *X_{j+1}'
        else:
            string_y = string_y + f' {round(abs(i[j]), 2)} *X_{j+1}'
    if k !=2 and k !=4:
        print(string_y)